#mask.py

import numpy as np
import pkg_resources
import matplotlib.pyplot as plt

def generate_mask(b):
    """
    生成掩膜版图
    
    参数：
    b (int): 传入的 b 值, 0 为不启用 OPC，1 为启用 OPC
    
    返回：
    fig (matplotlib.figure.Figure): Matplotlib 图形容器
    """
   # 加载文件
    file_path = pkg_resources.resource_filename(__name__, "mask.npy")
    amask = np.load(file_path)
    
    # 创建 Matplotlib 图形容器，祖传配方3:1
    fig, axs = plt.subplots(1, 5, figsize=(12, 4))

    # 绘制FUDAN五个字母对应的掩膜版
    for i in range(5):
        axs[i].imshow(amask[i + 5 * b], cmap='gray')
        axs[i].set_title(f'Mask {i}')

    # 调整子图布局
    plt.tight_layout()

    return fig  # 返回 Matplotlib 图形容器

# 掩膜版图函数调用示例
#plotter = generate_mask(0)
#plt.show()
